package com.szgz.platform.handler;

import com.jcraft.jsch.*;
import com.jcraft.jsch.ChannelShell;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * 终端WebSocket处理器
 * 处理WebSocket连接并建立SSH隧道到Docker容器
 */
@Component
public class TerminalWebSocketHandler implements WebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(TerminalWebSocketHandler.class);
    
    // 存储WebSocket会话和对应的SSH连接
    private final ConcurrentHashMap<String, SSHConnection> connections = new ConcurrentHashMap<>();
    private final ExecutorService executorService = Executors.newCachedThreadPool();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        logger.info("WebSocket连接建立: {}", session.getId());
        
        try {
            // 建立SSH连接到现有容器
            SSHConnection sshConnection = createSSHConnection();
            connections.put(session.getId(), sshConnection);
            
            // 启动数据转发线程
            startDataForwarding(session, sshConnection);
            
            session.sendMessage(new TextMessage("终端连接已建立"));
        } catch (Exception e) {
            logger.error("建立SSH连接失败", e);
            session.sendMessage(new TextMessage("连接失败: " + e.getMessage()));
            session.close();
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        if (message instanceof TextMessage) {
            String payload = ((TextMessage) message).getPayload();
            SSHConnection sshConnection = connections.get(session.getId());
            
            if (sshConnection != null && sshConnection.outputStream != null) {
                try {
                    sshConnection.outputStream.write(payload.getBytes());
                    sshConnection.outputStream.flush();
                } catch (IOException e) {
                    logger.error("发送数据到SSH失败", e);
                }
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        logger.error("WebSocket传输错误: {}", session.getId(), exception);
        closeConnection(session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        logger.info("WebSocket连接关闭: {}", session.getId());
        closeConnection(session.getId());
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 创建SSH连接到Docker容器
     */
    private SSHConnection createSSHConnection() throws Exception {
        JSch jsch = new JSch();
        Session session = jsch.getSession("root", "192.168.1.168", 2222);
        session.setPassword("root123"); // 使用root作为默认密码
        session.setConfig("StrictHostKeyChecking", "no");
        session.connect();
        
        ChannelShell channel = (ChannelShell) session.openChannel("shell");
        // 设置PTY（伪终端）以支持交互式终端
        channel.setPty(true);
        channel.setPtyType("xterm");
        channel.setPtySize(80, 24, 640, 480);
        channel.connect();
        
        return new SSHConnection(session, channel, 
            channel.getInputStream(), channel.getOutputStream());
    }

    /**
     * 启动数据转发线程
     */
    private void startDataForwarding(WebSocketSession webSocketSession, SSHConnection sshConnection) {
        executorService.submit(() -> {
            byte[] buffer = new byte[1024];
            try {
                while (webSocketSession.isOpen() && sshConnection.channel.isConnected()) {
                    int bytesRead = sshConnection.inputStream.read(buffer);
                    if (bytesRead > 0) {
                        String data = new String(buffer, 0, bytesRead);
                        // 同步发送消息，避免并发状态冲突
                        synchronized (webSocketSession) {
                            if (webSocketSession.isOpen()) {
                                webSocketSession.sendMessage(new TextMessage(data));
                            }
                        }
                    }
                }
            } catch (Exception e) {
                logger.error("数据转发异常", e);
            }
        });
    }

    /**
     * 关闭连接
     */
    private void closeConnection(String sessionId) {
        SSHConnection sshConnection = connections.remove(sessionId);
        if (sshConnection != null) {
            try {
                if (sshConnection.channel != null) {
                    sshConnection.channel.disconnect();
                }
                if (sshConnection.session != null) {
                    sshConnection.session.disconnect();
                }
            } catch (Exception e) {
                logger.error("关闭SSH连接失败", e);
            }
        }
    }

    /**
     * 获取查询参数
     */
    private String getQueryParam(WebSocketSession session, String paramName) {
        String query = session.getUri().getQuery();
        if (query != null) {
            String[] params = query.split("&");
            for (String param : params) {
                String[] keyValue = param.split("=");
                if (keyValue.length == 2 && keyValue[0].equals(paramName)) {
                    return keyValue[1];
                }
            }
        }
        return null;
    }

    /**
     * SSH连接信息封装类
     */
    private static class SSHConnection {
        final Session session;
        final Channel channel;
        final InputStream inputStream;
        final OutputStream outputStream;

        SSHConnection(Session session, Channel channel, InputStream inputStream, OutputStream outputStream) {
            this.session = session;
            this.channel = channel;
            this.inputStream = inputStream;
            this.outputStream = outputStream;
        }
    }
}